#计算评价指标,第二个函数calculate_trajectory_metrics还没有改好
import numpy as np
import torch
from sklearn.metrics import mean_squared_error, mean_absolute_error, mean_absolute_percentage_error

def calculate_batch_metrics(y_true_batch, y_pred_batch):

    # rmse = mean_squared_error(y_true_batch, y_pred_batch, squared=False)
    # mae = mean_absolute_error(y_true_batch, y_pred_batch)
    # mape = mean_absolute_percentage_error(y_true_batch, y_pred_batch)
    mse_per_component = torch.mean((y_true_batch - y_pred_batch) ** 2, dim=0)
    rmse_per_component = torch.sqrt(mse_per_component)
    # 计算 MAE（在每个分量上）
    mae_per_component = torch.mean(torch.abs(y_true_batch - y_pred_batch), dim=0)
    # 计算 MAPE（在每个分量上）
    mape_per_component = torch.mean(torch.abs((y_true_batch - y_pred_batch) / y_true_batch), dim=0) #除0会导致inf

    return rmse_per_component, mae_per_component, mape_per_component

 #sin里面的内容还没有改好！！！！
def calculate_trajectory_metrics(y_true, y_pred):
    T = len(y_true)
    
    EE = 0
    ATE = 0
    CTE = 0
    AE = 0
    
    for t in range(T):
        y_t_true = y_true[t]
        y_t_pred = y_pred[t]
        
        EE += np.sqrt(
            (y_t_true[0] - y_t_pred[0])**2 +
            (y_t_true[1] - y_t_pred[1])**2 +
            (y_t_true[2] - y_t_pred[2])**2
        )
        
        ATE += (
            (y_t_true[0] - y_t_pred[0]) * np.sin(y_t_true[3]) + 
            (y_t_true[1] - y_t_pred[1]) * np.cos(y_t_true[3])
        )
        
        CTE += (
            (y_t_true[0] - y_t_pred[0]) * np.cos(y_t_true[3]) -
            (y_t_true[1] - y_t_pred[1]) * np.sin(y_t_true[3])
        )
        
        AE += (y_t_true[2] - y_t_pred[2])**2
    
    EE = (1 / T) * EE
    ATE = (1 / T) * ATE
    CTE = (1 / T) * CTE
    AE = (1 / T) * np.sqrt(AE)
    
    return EE, ATE, CTE, AE